
#include <ctype.h>
#include <fcntl.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>

#include "bitfile.h"
#include "flash.h"

#define reg_read32(reg) (*((volatile uint32_t *)(reg)))
#define reg_write32(reg, val) (*((volatile uint32_t *)(reg))) = (val)


//===============================================================================
//二次开发仅仅需要关心以下1个变量
#define ALEXQSPI_ADDR 0x50000 
//===============================================================================

#define MAX_PATH    4096
#define MAX_SEGMENTS 8

uint8_t total_data[65536];

uint32_t reverse_bits_32(uint32_t x)
{
    x = ((x & 0x55555555) <<  1) | ((x & 0xAAAAAAAA) >>  1);
    x = ((x & 0x33333333) <<  2) | ((x & 0xCCCCCCCC) >>  2);
    x = ((x & 0x0F0F0F0F) <<  4) | ((x & 0xF0F0F0F0) >>  4);
    x = ((x & 0x00FF00FF) <<  8) | ((x & 0xFF00FF00) >>  8);
    x = ((x & 0x0000FFFF) << 16) | ((x & 0xFFFF0000) >> 16);
    return x;
}

uint16_t reverse_bits_16(uint16_t x)
{
    x = ((x & 0x5555) << 1) | ((x & 0xAAAA) >> 1);
    x = ((x & 0x3333) << 2) | ((x & 0xCCCC) >> 2);
    x = ((x & 0x0F0F) << 4) | ((x & 0xF0F0) >> 4);
    x = ((x & 0x00FF) << 8) | ((x & 0xFF00) >> 8);
    return x;
}

uint8_t reverse_bits_8(uint8_t x)
{
    x = ((x & 0x55) << 1) | ((x & 0xAA) >> 1);
    x = ((x & 0x33) << 2) | ((x & 0xCC) >> 2);
    x = ((x & 0x0F) << 4) | ((x & 0xF0) >> 4);
    return x;
}

char* stristr(const char *str1, const char *str2)
{
    const char* p1 = str1;
    const char* p2 = str2;
    const char* r = *p2 == 0 ? str1 : 0;

    while (*p1 != 0 && *p2 != 0)
    {
        if (tolower(*p1) == tolower(*p2))
        {
            if (r == 0)
            {
                r = p1;
            }

            p2++;
        }
        else
        {
            p2 = str2;
            if (r != 0)
            {
                p1 = r + 1;
            }

            if (tolower(*p1) == tolower(*p2))
            {
                r = p1;
                p2++;
            }
            else
            {
                r = 0;
            }
        }

        p1++;
    }

    return *p2 == 0 ? (char *)r : 0;
}

static void usage(char *name)
{
    fprintf(stderr,
        "usage: %s [options]\n"
        " -s slot    slot to program(选择需要操作的分区0:1)\n"
        " -r file    read flash to file(读数据,默认读128M)\n"
        " -w file    write and verify flash from file(烧写固件)\n"
        " -e         erase flash(擦除)\n"
        " -b         boot FPGA from flash(未实现)\n"
        " -t         hot reset FPGA(未实现)\n"
        " -y         no interactive confirm(不需要二次确认)\n",
        name);
}

/**
 * @description: 读取指定内存的数据
 * @param {flash_device} *fdev
 * @param {size_t} addr
 * @param {size_t} len
 * @param {void} *dest
 * @return {*}
 */
int flash_read_progress(struct flash_device *fdev, size_t addr, size_t len, void *dest)
{
    int ret = 0;
    size_t remain = len;
    size_t seg;
    int step = 0x10000;

    printf("Start address: 0x%08lx\n", addr);
    printf("Length: 0x%08lx\n", len);

    while (remain > 0)
    {
        if (remain > step)
        {
            // longer than step, trim
            if ((addr + step) & (step-1))
            {
                // align to step size
                seg = step - ((addr + step) & (step-1));
            }
            else
            {
                // already aligned
                seg = step;
            }
        }
        else
        {
            // shorter than step
            seg = remain;
        }

        printf("Read address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
        fflush(stdout);

        ret = flash_read(fdev, addr, seg, dest);

        if (ret) {
            fprintf(stderr, "\nRead failed\n");
            goto err;
        }

        addr += seg;
        remain -= seg;
        dest += seg;
    }

    printf("Read address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
    fflush(stdout);
    printf("\n");

err:
    return ret;
}

/**
 * @description: 不带数据校验的烧写
 * @param {flash_device} *fdev
 * @param {size_t} addr
 * @param {size_t} len
 * @param {void} *src
 * @return {*}
 */
int flash_write_progress(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
    int ret = 0;
    size_t remain = len;
    size_t seg;
    int step = 0x100; //页写

    printf("Start address: 0x%08lx\n", addr);
    printf("Length: 0x%08lx\n", len);

    step = fdev->write_buffer_size > step ? fdev->write_buffer_size : step;

    while (remain > 0)
    {
        if (remain > step)
        {
            // longer than step, trim
            if ((addr + step) & (step-1))
            {
                // align to step size
                seg = step - ((addr + step) & (step-1));
            }
            else
            {
                // already aligned
                seg = step;
            }
        }
        else
        {
            // shorter than step
            seg = remain;
        }

        printf("Write address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
        fflush(stdout);

        ret = flash_write(fdev, addr, seg, src);

        if (ret) {
            fprintf(stderr, "\nWrite failed\n");
            goto err;
        }

        addr += seg;
        remain -= seg;
        src += seg;
    }

    printf("Write address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
    fflush(stdout);
    printf("\n");

err:
    return ret;
}


/**
 * @description: 烧写程序带数据校验
 * @param {flash_device} *fdev
 * @param {size_t} addr
 * @param {size_t} len
 * @param {void} *src
 * @return {*}
 */
int flash_write_verify_progress(struct flash_device *fdev, size_t addr, size_t len, const void *src)
{
    int ret = 0;
    size_t remain = len;
    size_t seg;
    int step = 0x100;
    const uint8_t *ptr = src;
    uint8_t check_buf[step];
    memset(check_buf,0,step);

    printf("Start address: 0x%08lx\n", addr);
    printf("Length: 0x%08lx\n", len);

    step = fdev->write_buffer_size > step ? fdev->write_buffer_size : step;

    while (remain > 0)
    {
        if (remain > step)
        {
            // longer than step, trim
            if ((addr + step) & (step-1))
            {
                // align to step size
                seg = step - ((addr + step) & (step-1));
            }
            else
            {
                // already aligned
                seg = step;
            }
        }
        else
        {
            // shorter than step
            seg = remain;
        }

        printf("Write/verify address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
        fflush(stdout);

        ret = flash_write(fdev, addr, seg, ptr);

        if (ret) {
            fprintf(stderr, "\nWrite failed\n");
            goto err1;
        }

        for (int read_attempts = 2; read_attempts >= 0; read_attempts--)
        {
            ret = flash_read(fdev, addr, seg, check_buf);

            if (ret) {
                fprintf(stderr, "\nRead failed\n");
                goto err1;
            }

            if (memcmp(ptr, check_buf, seg))
            {
                fprintf(stderr, "\nVerify failed (%d more attempts)\n", read_attempts);
                
                for (size_t k = 0; k < seg; k++)
                {
                    if (ptr[k] != check_buf[k])
                    {
                        fprintf(stderr, "flash offset 0x%08lx: expected 0x%02x, read 0x%02x\n",
                            addr+k, ptr[k], check_buf[k]);
                    }
                }

                if (read_attempts > 0)
                    continue;

                ret = -1;

                printf("烧写失败，请重试...\r\n");
                printf("如果多次不成功，请联系管理员...\r\n");

                goto err1;
            }
        }

        addr += seg;
        remain -= seg;
        ptr += seg;
    }

    printf("Write/verify address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, (100*(len-remain))/len);
    fflush(stdout);
    printf("\n");

err1:
    
    return ret;
}

/**
 * @description: 擦除数据
 * @param {flash_device} *fdev
 * @param {size_t} addr
 * @param {size_t} len
 * @return {*}
 */
int flash_erase_progress(struct flash_device *fdev, size_t addr, size_t len)
{
    int ret;
    size_t remain = len;
    size_t seg;
    int step = 0x10000;

    printf("Start address: 0x%08lx\n", addr);
    printf("Length: 0x%08lx\n", len);

    step = fdev->erase_block_size > step ? fdev->erase_block_size : step;

    while (remain > 0)
    {
        if (remain > step)
        {
            // longer than step, trim
            if ((addr + step) & (step-1))
            {
                // align to step size
                seg = step - ((addr + step) & (step-1));
            }
            else
            {
                // already aligned
                seg = step;
            }
        }
        else
        {
            // shorter than step
            seg = remain;
        }

        printf("Erase address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, ((100*(len-remain))/len));
        fflush(stdout);

        ret = flash_erase(fdev, addr, seg);

        if (ret)
            return ret;

        addr += seg;
        remain -= seg;
    }

    printf("Erase address 0x%08lx, length 0x%08lx (%ld%%)\r", addr, seg, ((100*(len-remain))/len));
    fflush(stdout);
    printf("\n");

    return 0;
}

int write_str_to_file(const char *file_name, const char *str)
{
    int ret = 0;
    FILE *fp = fopen(file_name, "w");

    if (!fp)
    {
        perror("failed to open file");
        return -1;
    }

    if (fputs(str, fp) == EOF)
    {
        perror("failed to write to file");
        ret = -1;
    }

    fclose(fp);
    return ret;
}

int write_1_to_file(const char *file_name)
{
    return write_str_to_file(file_name, "1");
}

#define FILE_TYPE_BIN 0
#define FILE_TYPE_HEX 1
#define FILE_TYPE_BIT 2

int file_type_from_ext(const char *file_name)
{
    char *ptr;
    char buffer[32];

    ptr = strrchr(file_name, '.');

    if (!ptr)
    {
        return FILE_TYPE_BIN;
    }

    ptr++;

    for (int i = 0; i < sizeof(buffer)-1 && *ptr; i++)
    {
        buffer[i] = tolower(*ptr++);
        buffer[i+1] = 0;
    }

    if (strcmp(buffer, "hex") == 0 || strcmp(buffer, "mcs") == 0)
    {
        return FILE_TYPE_HEX;
    }

    if (strcmp(buffer, "bit") == 0)
    {
        return FILE_TYPE_BIT;
    }

    return FILE_TYPE_BIN;
}


/**
 * @description: 主函数
 * @param {int} argc
 * @param {char} *argv
 * @return {*}
 */
int main(int argc, char *argv[])
{
    char *name;
    int opt;
    int ret = 0;

    char *read_file_name = NULL;
    FILE *read_file = NULL;
    char *write_file_name = NULL;
    FILE *write_file = NULL;


    int slot = -1;

    //指令模式使能标志
    char action_read = 0;
    char action_write = 0;
    char action_erase = 0;
    char action_boot = 0;
    char action_reset = 0;
    char no_confirm = 0;


    struct flash_device *pri_flash = NULL;

    int flash_segment_count = 0;
    size_t flash_segment_start[MAX_SEGMENTS];
    size_t flash_segment_length[MAX_SEGMENTS];

    name = strrchr(argv[0], '/');
    name = name ? 1+name : argv[0];

    while ((opt = getopt(argc, argv, "s:r:w:ebtyh?")) != EOF)
    {
        switch (opt)
        {
        case 's':
            slot = atoi(optarg);
            break;
        case 'r':
            action_read = 1;
            read_file_name = optarg;
            break;
        case 'w':
            action_write = 1;
            write_file_name = optarg;
            break;
        case 'e':
            action_erase = 1;
            break;
        case 'b':
            action_boot = 1;
            action_reset = 1;
            break;
        case 't':
            action_reset = 1;
            break;
        case 'y':
            no_confirm = 1;
            break;
        case 'h':
        case '?':
            usage(name);
            return 0;
        default:
            usage(name);
            return -1;
        }
    }

    uint8_t flash_configuration = 0;
    uint8_t flash_data_width = 0;       //QSPI默认8位宽
    uint8_t flash_default_segment = 1;  //FLASH烧写的默认分区
    uint8_t flash_fallback_segment = 0; //FLASH的备用分区
    uint32_t flash_segment0_length = 0; //FLASH第一个分区的大小

    int bitswap = 0;
    int word_size = 8;

    size_t flash_size = 0;
    size_t segment_size = 0;
    size_t segment_offset = 0;
    flash_data_width = 8;   

    pri_flash = flash_open_qspi(flash_data_width);// flash_data_width==4

    if (!pri_flash)
    {
        fprintf(stderr, "Failed to connect to flash device\n");
        ret = -1;
        goto err;
    }

    flash_size = pri_flash -> size;//这个size是从Flash内部的寄存器读回的！在flash_open_spi内部的init函数里操作

    flash_configuration    = 2;
    flash_default_segment  = 1;
    flash_fallback_segment = 0;
    flash_segment0_length  = 0x08000000;
		

    switch (flash_configuration)
    {
        case 0:
        case 1:
            flash_segment_count = 1;
            flash_segment_start[0] = 0;
            flash_segment_length[0] = flash_size;
            break;
        case 2:
            if (flash_segment0_length == 0)
            {
                flash_segment0_length = flash_size >> 1;
            }
            else if (flash_size < flash_segment0_length)
            {
                fprintf(stderr, "Invalid flash configuration\n");
                ret = -1;
                goto err;
            }

            flash_segment_count = 2;
            flash_segment_start[0] = 0;
            flash_segment_length[0] = flash_segment0_length;
            flash_segment_start[1] = flash_segment_start[0]+flash_segment_length[0];
            flash_segment_length[1] = flash_size-flash_segment_start[1];
            break;
        case 4:
            flash_segment_count = 4;
            flash_segment_start[0] = 0;
            flash_segment_length[0] = flash_size >> 2;
            for (int k = 1; k < 4; k++)
            {
                flash_segment_start[k] = flash_segment_start[k-1]+flash_segment_length[k-1];
                flash_segment_length[k] = flash_size >> 2;
            }
            break;
        case 8:
            flash_segment_count = 8;
            flash_segment_start[0] = 0;
            flash_segment_length[0] = flash_size >> 3;
            for (int k = 1; k < 8; k++)
            {
                flash_segment_start[k] = flash_segment_start[k-1]+flash_segment_length[k-1];
                flash_segment_length[k] = flash_size >> 3;
            }
            break;
        default:
            fprintf(stderr, "Unknown flash configuration (0x%02x)\n", flash_configuration);
            ret = -1;
            goto err;
    }

    for (int k = 0; k < flash_segment_count; k++)
    {
        printf("Flash segment %d: start 0x%08lx length 0x%08lx\n", k, flash_segment_start[k], flash_segment_length[k]);
    }

    printf("Default segment: %d\n", flash_default_segment);
    if (flash_fallback_segment == flash_default_segment || flash_fallback_segment >= flash_segment_count)
    {
        printf("Fallback segment: none\n");
    }
    else
    {
        printf("Fallback segment: %d\n", flash_fallback_segment);
    }

    if (slot < 0)
    {
        slot = flash_default_segment;
    }

    if ((action_read || action_write) && (slot < 0 || slot >= flash_segment_count))
    {
        fprintf(stderr, "Requested slot is not valid (%d)\n", slot);
        ret = -1;
        goto err;
    }

    segment_offset = flash_segment_start[slot];
    segment_size = flash_segment_length[slot];

    printf("Selected: segment %d start 0x%08lx length 0x%08lx\n", slot, segment_offset, segment_size);

    if (action_erase)
    {
        if (!no_confirm)
        {
            char str[32];

            printf("Are you sure you want to erase the selected segment?\n");
            printf("[y/N]: ");

            fgets(str, sizeof(str), stdin);

            if (str[0] != 'y' && str[0] != 'Y')
                goto err;
        }

        printf("Erasing flash...\n");
        if (flash_erase_progress(pri_flash, segment_offset, segment_size))
        {
            fprintf(stderr, "Erase failed!\n");
            ret = -1;
            goto err;
        }

        printf("Erase complete!\n");
        
    }

    if (action_write)
    {
        char *segment = calloc(segment_size, 1);
        memset(segment, 0xff, segment_size);
        size_t len;

        int file_type = file_type_from_ext(write_file_name);

        if (file_type == FILE_TYPE_BIN)
        {
            // read binary file
            printf("Reading binary file \"%s\"...\n", write_file_name);
            write_file = fopen(write_file_name, "rb");

            if (!write_file)
            {
                fprintf(stderr, "Failed to open file\n");
                free(segment);
                ret = -1;
                goto err;
            }

            fseek(write_file, 0, SEEK_END);
            len = ftell(write_file);
            rewind(write_file);

            if (len > segment_size)
            {
                fprintf(stderr, "File larger than segment (%ld > %ld)\n", len, segment_size);
                fclose(write_file);
                free(segment);
                ret = -1;
                goto err;
            }

            if (fread(segment, 1, len, write_file) < len)
            {
                fprintf(stderr, "Error reading file\n");
                fclose(write_file);
                free(segment);
                ret = -1;
                goto err;
            }

            fclose(write_file);
        }
        else if (file_type == FILE_TYPE_BIT)
        {
            // read bit file
            struct bitfile *bf;

            bf = bitfile_create_from_file(write_file_name);

            if (!bf)
            {
                fprintf(stderr, "Error reading bit file\n");
                free(segment);
                ret = -1;
                goto err;
            }

            if (bf->data_len > segment_size)
            {
                fprintf(stderr, "File larger than segment (%ld > %ld)\n", bf->data_len, segment_size);
                bitfile_close(bf);
                free(segment);
                ret = -1;
                goto err;
            }

            len = bf->data_len;
            memcpy(segment, bf->data, bf->data_len);

            bitfile_close(bf);
        }
        else if (file_type == FILE_TYPE_HEX)
        {
            fprintf(stderr, "Hex files are not currently supported\n");
            free(segment);
            ret = -1;
            goto err;
        }
        else
        {
            fprintf(stderr, "Unsupported file type\n");
            free(segment);
            ret = -1;
            goto err;
        }

        // check sync word
        if (memcmp(segment+0x50, "\xAA\x99\x55\x66", 4))
        {
            fprintf(stderr, "Bitstream sync word not found\n");
            free(segment);
            ret = -1;
            goto err;
        }

        if (bitswap)
        {
            if (word_size == 16)
            {
                uint16_t *p = (uint16_t *)segment;

                for (size_t k = 0; k < segment_size; k += 2)
                {
                    *p = reverse_bits_16(*p);
                    p++;
                }
            }
            else
            {
                uint8_t *p = (uint8_t *)segment;

                for (size_t k = 0; k < segment_size; k++)
                {
                    *p = reverse_bits_8(*p);
                    p++;
                }
            }
        }

        // round up length to block size
        if ((segment_offset + len) & (pri_flash->erase_block_size-1))
        {
            len += pri_flash->erase_block_size - ((segment_offset + len) & (pri_flash->erase_block_size-1));
        }

        if (!no_confirm)
        {
            char str[32];

            printf("Are you sure you want to write the selected segment?\n");
            printf("[y/N]: ");

            fgets(str, sizeof(str), stdin);

            if (str[0] != 'y' && str[0] != 'Y')
                goto err;
        }

        printf("Erasing flash...\n");
        if (flash_erase_progress(pri_flash, segment_offset, len))
        {
            fprintf(stderr, "Erase failed!\n");
            ret = -1;
            free(segment);
            goto err;
        }

        printf("Writing and verifying flash...\n");
        if (flash_write_verify_progress(pri_flash, segment_offset, len, segment))
        {
            fprintf(stderr, "Write/verify failed!\n");
            ret = -1;
            free(segment);
            goto err;
        }
        //  预留无校验烧写
        // printf("Writing flash...\n");
        // if (flash_write_progress(pri_flash, segment_offset, len, segment))
        // {
        //     fprintf(stderr, "Write failed!\n");
        //     ret = -1;
        //     free(segment);
        //     goto err;
        // }

        printf("Programming succeeded!\n");
        

        free(segment);
    }

    if (action_read)
    {
        char *segment = calloc(segment_size, 1);
        memset(segment, 0xff, segment_size);

        printf("Reading flash...\n");
        flash_read_progress(pri_flash, segment_offset, segment_size, segment);
        

        if (bitswap)
        {
            if (word_size == 16)
            {
                uint16_t *p = (uint16_t *)segment;

                for (size_t k = 0; k < segment_size; k += 2)
                {
                    *p = reverse_bits_16(*p);
                    p++;
                }
            }
            else
            {
                uint8_t *p = (uint8_t *)segment;

                for (size_t k = 0; k < segment_size; k++)
                {
                    *p = reverse_bits_8(*p);
                    p++;
                }
            }
        }

        int file_type = file_type_from_ext(read_file_name);

        if (file_type == FILE_TYPE_BIN)
        {
            // write binary file
            printf("Writing binary file \"%s\"...\n", read_file_name);
            read_file = fopen(read_file_name, "wb");
            fwrite(segment, 1, segment_size, read_file);
            fclose(read_file);
        }
        else if (file_type == FILE_TYPE_HEX)
        {
            fprintf(stderr, "Hex files are not currently supported\n");
            free(segment);
            ret = -1;
            goto err;
        }
        else
        {
            fprintf(stderr, "Unsupported file type\n");
            free(segment);
            ret = -1;
            goto err;
        }

        free(segment);
    }

err:

    flash_release(pri_flash);
    pri_flash = NULL;

    return ret;

}




